from vBaum import *
from math import *

from tree_common import *

# Monopodial: Main axis dominates
class SympodialBranch:
	def __init__(self):
		self.branchAngle = radians(60)
		self.branchAngleSigma = radians(20)
		
		#numBranches depending on the evaluation of the random function:
		#0:  no branch
		#1:  monochasium (dominant branch + recessive branch)
		#    branchAngle will affect the recessive branch.
		#    the dominant branch will be adjusted according to the
		#    branch angle rule
		#2:  dichasium (equal angles for each branch, same diameters)
		#3+: pleiochasium (same as dichasium, more branches)
		self.numBranchesPerNode = 1
		self.numBranchesPerNodeRange = 0.7 #sigma * numBranchesPerNode
		
		self.angleBetweenBranchesAtNodeSigma = radians(10)
		self.divergence = True
		self.successor = self
		self.orientUp = False
		self.radiusFallOff = 0.5 # 0.5

		self.roundNodes = True
		self.nodeRandomness = 0.05 #per one unit branch length
		self.branchRoughness = 0.6
		self.initialInterNodeLength = 0.15
		self.interNodeLengthGrowth = 0.04
		self.interNodePoints = 3 # sets number of spline points on internodes

		self.interNodeAreaLoss = 0.1
		self.angleOffsetSigma = 0.1
		self.minBranchDiameter = voxelSize()
		self.minBranchLength = max(self.initialInterNodeLength*0.75, voxelSize()*7)
		
		self.minRecessiveAreaRatio = 0.1
		self.recessiveAreaMultiplier = 1.5
		
		self.branchRadiusScale = 1.0

		#used to calculate RadiusHeightConstant
		designLength = 5
		designRadius = 0.25
		
		self.RadiusHeightConstant = calcRadiusLengthConstant(designLength, designRadius, self.radiusFallOff)
		
	def heightToDiam(self, length):
		return LengthToDiam(length, self.RadiusHeightConstant, self.radiusFallOff)
				
	def generateGeometry(self, startPos, direction, trunkRadius):
		#for graphics output, don't generate small diameters
		graphicsTrunkRadius = max(self.minBranchDiameter, trunkRadius)
		
		#penalize small branches a bit
		maxLength = diamToLength(trunkRadius, self.RadiusHeightConstant, self.radiusFallOff)

		#reject too small branches
		if(maxLength < self.minBranchLength):
			return Volume()
		
		print("Max length is ", maxLength)
				
		endPos = startPos + direction * maxLength
		
		nodeHeights = generateNodePositions(maxLength, self.initialInterNodeLength, self.interNodeLengthGrowth)
		
		#shorten branch so node positions are ideal
#		if len(nodeHeights) > 1:
#			shortenAmount = nodeHeights[0]
#			for i in range(len(nodeHeights)):
#				nodeHeights[i] = nodeHeights[i]-shortenAmount
		
		print("length of nodeHeights is", len(nodeHeights))
		print("nodeHeights:", nodeHeights)

		actualLength = nodeHeights[-1]
		print("actual length is ", actualLength)
		
		print("trunk radius:", trunkRadius)
		
		#angles and radii of branches
		#first entry is always the main branch
	
		aaBranchAngles = []
		aaBranchRadii = []
		bPleioEnding = False

		angleOffset = 0
		
		if len(nodeHeights) > 1:
			#small branch
			currRadius = (1 - nodeHeights[1] / actualLength) * trunkRadius
			currArea = currRadius * currRadius
		
		#don't need branches at the very end of this branch
		for i in range(1, len(nodeHeights)-1):
			print("curr node height:", nodeHeights[i])

			numBranches = rand.gauss(0, self.numBranchesPerNode * self.numBranchesPerNodeRange) + self.numBranchesPerNode	
			numBranches = max(round(numBranches), 0)
			
			#HACK force monochasium / no branch
			numBranches = min(1, numBranches)
			
			print("generating", numBranches, "branches.")
			
			branchAngle = rand.gauss(self.branchAngle, self.branchAngleSigma)
			
			branchAngle = min(max(-90, branchAngle), radians(90))
			
			branchAngles = []
			branchRatios = []

			if numBranches == 0:
				#no branch
				branchAngles = [0]
				branchRatios = [1 - self.interNodeAreaLoss]
				
			elif numBranches == 1:
				#monochasium
				#calculate how much of the cross section goes into the recessive branch
				recessiveAreaRatio = 0.5 * min(1, max(self.minRecessiveAreaRatio, cos(branchAngle) + self.minRecessiveAreaRatio)*self.recessiveAreaMultiplier)
				
				dominantAreaRatio = 1 - recessiveAreaRatio
				
				#calculate branch angle of the dominant branch
				#dominantBranchAngle = acos(dominantAreaRatio)
#				dominantBranchAngle = recessiveAreaRatio * 2.0 * branchAngle
				dominantBranchAngle = recessiveAreaRatio * branchAngle
				
				branchAngles = [dominantBranchAngle, branchAngle]
				branchRatios = [dominantAreaRatio, recessiveAreaRatio]
			else:
				#di- or pleiochasium
				bPleioEnding = True
				branchAngles = [branchAngle]*numBranches
				branchRatios = [1.0 / numBranches]*numBranches

			branchRadii = []
			for branchRatio in branchRatios:
				branchRadii.append( sqrt(branchRatio * currArea) )
				
			currArea = currArea * branchRatios[0]
				
			aaBranchAngles.append(branchAngles)
			aaBranchRadii.append(branchRadii)
			
			if bPleioEnding:
				break;
						
		print("finished creating branch angles")
		for branchAngles in aaBranchAngles:
			print("node")
			for angle in branchAngles:
				print(degrees(angle))

		print("radii:", aaBranchRadii)
		
		#generate branches using angles and radii
		
		#iterate until there is a pleiochasium node
		#in that case, we have to start a new branch by recursion
		nodeDirection = direction
		
		nodePositions = [startPos, startPos + (nodeDirection * nodeHeights[0])]
		nodeDirections = [nodeDirection]
		
		if len(nodeHeights) > 1:
			#omit the first generated node
			nodePositions.append(startPos + (nodeDirection * nodeHeights[1]))
			
		nodePos = nodePositions[-1]
			
		aFirstBranchAxialAngle = []
		
		for i in range(0, len(aaBranchAngles)-1):
			nodeHeight = nodeHeights[1 + i];
			interNodeLen = nodeHeights[2 + i] - nodeHeight
			print("generating for height", nodeHeight, ", internode length", interNodeLen)
			
			#if self.divergence: angle = angleOffset
			
			print("current node position:")
			printVec(nodePos)
			
			#randomBranchAngle = rand.gauss(aaBranchAngles[i][0], self.branchAngleSigma)
			randomBranchAngle = aaBranchAngles[i][0]
			
			print("actual branch angle:", degrees(randomBranchAngle))
			
			firstBranchDir = computeFirstBranchDir(nodeDirection, randomBranchAngle)
			
			if self.orientUp:
				firstBranchDir = rotateDir(nodeDirection, Vec3(0, 1, 0), randomBranchAngle)
				
			angle = rand.random() * 2 * pi
			aFirstBranchAxialAngle.append(angle)
			
			print("axial angle", degrees(angle))
			
			branchDir = rotateDir(firstBranchDir, nodeDirection, rand.gauss(angle, self.angleBetweenBranchesAtNodeSigma))
			
			print("current dominant branch direction:")
			printVec(branchDir)
			
			newPos = branchDir * (interNodeLen) + nodePos
			nodePositions.append(newPos)
			
			nodeDirection = newPos - nodePos
			nodeDirection.normalise()
			
			nodeDirections.append(nodeDirection)
			
			nodePos = newPos

		
		#last entry is either pleio or end of tree
		#if it was end of tree, do nothing
		#if it was pleio, delete the last vertex and generate a pleio
		if bPleioEnding:
			nodePositions.pop()
			aaBranchAngles.pop()
			
		#randomize node positions
		nodePositions, branchContour = randomizeNodesOnBranch(nodePositions, nodeHeights, self.interNodePoints, self.branchRoughness, self.nodeRandomness)
		
		#create this branch
		print("branchPositions:")
		for position in branchContour:
			printVec(position)

		branch = createBranch(branchContour, graphicsTrunkRadius * self.branchRadiusScale, voxelSize())
		#create monochasium branches
		for i in range(0, len(aaBranchAngles)-1):
			if len(aaBranchAngles[i]) != 2:
				continue
		
			print("creating branch", i)
#			print("len(aaBranchAngles)", len(aaBranchAngles))
#			print("len(aaBranchRadii)", len(aaBranchRadii))
			
			nodeDirection = nodeDirections[i+1]
			
			#branchAngle = rand.gauss(aaBranchAngles[i][0], self.branchAngleSigma)
			branchAngle = aaBranchAngles[i][1]
			
			firstBranchDir = computeFirstBranchDir(nodeDirection, branchAngle)
			
			if self.orientUp:
				firstBranchDir = rotateDir(nodeDirection, Vec3(0, 1, 0), branchAngle)

			#load branch angle 
			#random was already added, so this angle is final
			antiAngle = aFirstBranchAxialAngle[i] + pi
			if antiAngle > 2*pi:
				antiAngle = antiAngle - 2*pi		

			branchDir = rotateDir(firstBranchDir, nodeDirection, antiAngle)
			branchPos = nodePositions[i+1]
			branchRadius = aaBranchRadii[i][1]
			
			recessiveSubBranch = self.generateGeometry(branchPos, branchDir, branchRadius)
			branch = branch.union(recessiveSubBranch)
			
#			return branch
			
#		if self.roundNodes:
#			branch = roundSubVolume(branch, nodePos, nodeRadius*2.0, nodeRadius*0.5)			
				
		return branch

node = SympodialBranch()
tree = node.generateGeometry(Vec3(0, 0, 0), Vec3(0, 1, 0), node.heightToDiam(7))
tree.exportAsObj("sympodialTree", voxelsPerUnit, 2)

